#include <ros/ros.h>
#include <sensor_msgs/Image.h>
#include <cv_bridge/cv_bridge.h>
#include <image_transport/image_transport.h>
#include <std_msgs/String.h>
#include "color_correction/color_correction.hpp"
#include <opencv2/opencv.hpp>
#include <opencv2/imgcodecs.hpp>
#include <map>
#include <string>
#include <locale>
#include <sys/stat.h>
#include <chrono>

class ColorCorrectionProcessor {
private:
    ros::NodeHandle nh_;
    ros::NodeHandle pnh_;
    image_transport::ImageTransport it_;
    
    // Topic mapping: input_topic -> {corrector, publisher}
    struct TopicHandler {
        std::unique_ptr<ColorCorrector> corrector;
        image_transport::Publisher publisher;
        image_transport::Subscriber subscriber;
    };
    
    std::map<std::string, std::unique_ptr<TopicHandler>> topic_handlers_;
    
    // Parameters
    std::string config_file_path_;
    std::vector<std::string> input_topics_;
    
    // File monitoring
    time_t last_config_mtime_;
    ros::Timer config_check_timer_;
    
public:
    ColorCorrectionProcessor() : nh_(), pnh_("~"), it_(nh_), last_config_mtime_(0) {
        loadParameters();
        initializeTopicHandlers();
        
        // Start file monitoring timer
        config_check_timer_ = nh_.createTimer(ros::Duration(1.0), 
            &ColorCorrectionProcessor::configFileCheckCallback, this);
        
        // Initialize config file modification time
        updateConfigFileTime();
        
        ROS_INFO("色彩校正处理器初始化完成");
        ROS_INFO("配置文件加载路径: %s", config_file_path_.c_str());
        ROS_INFO("处理 %zu 个输入话题", input_topics_.size());
        ROS_INFO("已启动配置文件监控 (1秒检查间隔)");
    }
    
private:
    void loadParameters() {
        pnh_.param<std::string>("config_file_path", config_file_path_, 
                               "/home/firefly/ws_vision/src/color_correction/config/color_correction_params.json");
        
        // Get list of input topics to process
        XmlRpc::XmlRpcValue topic_list;
        if (pnh_.getParam("input_topics", topic_list)) {
            if (topic_list.getType() == XmlRpc::XmlRpcValue::TypeArray) {
                for (int i = 0; i < topic_list.size(); ++i) {
                    if (topic_list[i].getType() == XmlRpc::XmlRpcValue::TypeString) {
                        input_topics_.push_back(static_cast<std::string>(topic_list[i]));
                    }
                }
            }
        }
        
        // Default topics if none specified
        if (input_topics_.empty()) {
            input_topics_.push_back("/camera/image_raw");
            input_topics_.push_back("/camera/image_color");
        }
        
        ROS_INFO("参数加载完成:");
        ROS_INFO("  配置文件路径: %s", config_file_path_.c_str());
        ROS_INFO("  输入话题数量: %zu 个", input_topics_.size());
        for (const auto& topic : input_topics_) {
            ROS_INFO("    - %s", topic.c_str());
        }
    }
    
    void initializeTopicHandlers() {
        for (const auto& input_topic : input_topics_) {
            auto handler = std::make_unique<TopicHandler>();
            
            // Create color corrector instance
            handler->corrector = std::make_unique<ColorCorrector>();
            
            // Load calibration parameters
            if (!handler->corrector->loadParametersFromJson(config_file_path_)) {
                ROS_WARN("话题 %s 的色彩校正参数加载失败", input_topic.c_str());
                ROS_WARN("将直接传递图像，不进行校正");
            } else {
                ROS_INFO("话题 %s 的色彩校正参数加载成功", input_topic.c_str());
            }
            
            // Create output topic name by adding _corrected suffix
            std::string output_topic = input_topic + "_corrected";
            
            // Create publisher
            handler->publisher = it_.advertise(output_topic, 1);
            
            // Create subscriber
            handler->subscriber = it_.subscribe(input_topic, 1, 
                boost::bind(&ColorCorrectionProcessor::imageCallback, this, _1, input_topic));
            
            // Store handler
            topic_handlers_[input_topic] = std::move(handler);
            
            ROS_INFO("话题设置完成: %s -> %s", input_topic.c_str(), output_topic.c_str());
        }
    }
    
    void imageCallback(const sensor_msgs::ImageConstPtr& msg, const std::string& input_topic) {
        auto it = topic_handlers_.find(input_topic);
        if (it == topic_handlers_.end()) {
            ROS_ERROR("未找到话题 %s 的处理器", input_topic.c_str());
            return;
        }
        
        auto& handler = it->second;
        
        try {
            // Convert ROS image to OpenCV image
            cv_bridge::CvImagePtr cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
            cv::Mat input_image = cv_ptr->image;
            
            cv::Mat output_image = input_image;
            
            // Apply color correction if calibrated
            if (handler->corrector->isCalibrationReady()) {
                output_image = handler->corrector->applyColorCorrection(input_image);
            } else {
                // Just pass through the original image
                ROS_DEBUG_THROTTLE(5.0, "话题 %s 的色彩校正器未校准，直接传递图像", 
                                  input_topic.c_str());
            }
            
            // Convert back to ROS image and publish
            cv_bridge::CvImage output_msg;
            output_msg.header = msg->header;
            output_msg.encoding = sensor_msgs::image_encodings::BGR8;
            output_msg.image = output_image;
            
            handler->publisher.publish(output_msg.toImageMsg());
            
        } catch (cv_bridge::Exception& e) {
            ROS_ERROR("话题 %s 的cv_bridge异常: %s", input_topic.c_str(), e.what());
        } catch (const std::exception& e) {
            ROS_ERROR("话题 %s 的色彩校正异常: %s", input_topic.c_str(), e.what());
        }
    }
    
public:
    void configFileCheckCallback(const ros::TimerEvent& event) {
        time_t current_mtime = getConfigFileModificationTime();
        
        if (current_mtime > last_config_mtime_ && last_config_mtime_ != 0) {
            ROS_INFO("检测到配置文件更新，正在重新加载参数...");
            reloadConfiguration();
            updateConfigFileTime();
        }
    }
    
    time_t getConfigFileModificationTime() {
        struct stat file_stat;
        if (stat(config_file_path_.c_str(), &file_stat) == 0) {
            return file_stat.st_mtime;
        }
        return 0;
    }
    
    void updateConfigFileTime() {
        last_config_mtime_ = getConfigFileModificationTime();
    }
    
    void reloadConfiguration() {
        ROS_INFO("正在重新加载色彩校正配置...");
        
        for (auto& pair : topic_handlers_) {
            const std::string& topic = pair.first;
            auto& handler = pair.second;
            
            if (handler->corrector->loadParametersFromJson(config_file_path_)) {
                ROS_INFO("话题 %s 的色彩校正参数重新加载成功", topic.c_str());
            } else {
                ROS_WARN("话题 %s 的色彩校正参数重新加载失败", topic.c_str());
            }
        }
    }
    
    void spin() {
        ros::spin();
    }
};

int main(int argc, char** argv) {
    std::setlocale(LC_ALL, "zh_CN.UTF-8");
    ros::init(argc, argv, "color_correction_processor");
    
    try {
        ColorCorrectionProcessor processor;
        processor.spin();
    } catch (const std::exception& e) {
        ROS_FATAL("色彩校正处理器失败: %s", e.what());
        return -1;
    }
    
    return 0;
}
